import xalpha as xa
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
# import seaborn as sns

from model import Model


MEAN = 8.676183e-05
STD = 0.031644963

@torch.no_grad()
def main(steps=10, n_sample=65536, replay=False):
    rbk = steps if replay else 0
    raw_data = xa.indexinfo('SH000300', save=True, path="data.csv")
    print(raw_data.price)

    rt = xa.universal.get_rt('SH000300')
    if rbk == 0 and str(raw_data.price['date'].iloc[-1])[:10] != rt["time"][:10]:
        patch = pd.DataFrame([[rt["time"], 0, rt["current"], 0]], columns=raw_data.price.columns)
        raw_data.price = pd.concat([raw_data.price, patch])

    if rbk:
        raw_price = raw_data.price.iloc[:-rbk]
        gt_price = raw_data.price.iloc[-rbk - 1:]
    else:
        raw_price = raw_data.price
    x_data = raw_price["totvalue"].to_numpy()

    x_data = np.log(x_data, dtype=np.float32)[-600:]
    x_data = x_data[1:] - x_data[:-1]
    x_data = (x_data - MEAN) / STD
    x_data = torch.from_numpy(x_data)[None, :, None]

    m = Model()
    m.load_state_dict(torch.load("300.pt", map_location='cpu'))
    dist, h = m(x_data[:, :-1], ret_h=True)
    h = (h[0].repeat(1, n_sample, 1), h[1].repeat(1, n_sample, 1))
    dist, h = m(x_data[:, -1:].repeat(n_sample, 1, 1), h, ret_h=True)

    # plt.plot(range(-steps, 0), x_data[0, -steps:, 0])

    samples = []
    _prev = np.float32(raw_price["totvalue"].iloc[-1])
    totvalue = torch.tensor(_prev).repeat(n_sample)
    for step in range(steps):
        x_pred = dist.sample()[..., None]
        totvalue *= (x_pred.flatten() * STD + MEAN).exp()
        samples.append(totvalue.clone())
        dist, h = m(x_pred, h, ret_h=True)
    samples = torch.stack(samples)
    mean = samples.mean(1)
    ci0 = np.percentile(samples, 95, 1)
    ci1 = np.percentile(samples, 5, 1)
    ci2 = np.percentile(samples, 75, 1)
    ci3 = np.percentile(samples, 25, 1)
    # plt.yscale('log')
    plt.plot(range(-1, steps), [_prev, *mean], label="$\mathbb{E}(\cdot)$")
    if n_sample > 1:
        plt.plot(range(-1, steps), [_prev, *ci0], c='gray', ls=('dotted'), label=".9CI")
        plt.plot(range(-1, steps), [_prev, *ci1], c='gray', ls=('dotted'))
        plt.plot(range(-1, steps), [_prev, *ci2], c='lightblue', ls=('dotted'), label=".5CI")
        plt.plot(range(-1, steps), [_prev, *ci3], c='lightblue', ls=('dotted'))
    plt.legend()

    # prev
    plt.plot(range(-steps // 2, 0), raw_price["totvalue"][-steps // 2:])

    # gt
    if rbk:
        plt.plot(range(-1, rbk), gt_price["totvalue"])
    plt.title(f"Expected annualized return: {((mean[-1].item() / _prev) ** (200 / steps) - 1) * 100: .2f}%")
    plt.xlabel("Days")
    plt.ylabel("SH000300")
    plt.savefig("plt.jpg")
    plt.close()

if __name__ == "__main__":
    main()
